package cn.demoncat.util.mybatis.util;

import java.util.concurrent.atomic.AtomicInteger;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import cn.demoncat.util.lang.RandomUtil;
import cn.demoncat.util.mybatis.constant.DataSourceLogConstant;
import cn.demoncat.util.mybatis.constant.DataSourceTypeConstant;

/**
 * 读写分离的路由数据源
 * 
 * @author 延晓磊
 *
 * @since 2017年10月13日
 */
public class RwRoutingDataSource extends AbstractRoutingDataSource {
	//从库数量
	private final int readDataSourceNum;
	//递增器
	private final AtomicInteger count = new AtomicInteger(0);
	//负载均衡类型：0 随机，1轮循
	private final int balanceType;
	
	private static final Logger LOG = LoggerFactory.getLogger(RwRoutingDataSource.class);
	
	/**
	 * 默认构造
	 * readDataSourceNum = 1
	 * balanceType = 0
	 */
	public RwRoutingDataSource(){
		this.readDataSourceNum = 1;
		this.balanceType = 0;
	}
	
	/**
	 * 构造器
	 * 指定readDataSourceNum和balanceType
	 */
	public RwRoutingDataSource(int readDataSourceNum,int balanceType) {
		this.readDataSourceNum = readDataSourceNum;
		this.balanceType = balanceType;
    }
	
	/**
	 * 确定当前数据源
	 * 
	 * @return
	 *
	 * @author 延晓磊
	 *
	 * @since 2017年10月13日
	 *
	 */
	@Override
	protected Object determineCurrentLookupKey() {
		String curr;
		if (StringUtils.isBlank(DataSourceHolder.get())) {
			//不设置时默认写库
			curr = DataSourceTypeConstant.WRITE;
		}else if (DataSourceHolder.get().equals(DataSourceTypeConstant.WRITE)){
			//写库
			curr = DataSourceTypeConstant.WRITE;
		}else{
			//读库
			switch (readDataSourceNum) {
			case 0:
				//没有读库
				curr = DataSourceTypeConstant.WRITE;
				break;
			case 1:
				//单读库
				curr = DataSourceTypeConstant.READ+0;
				break;
			default:
				if (balanceType == 0) {
					//随机负载均衡
					curr = DataSourceTypeConstant.READ + RandomUtil.number(0, readDataSourceNum);
				}else {
					//轮循负载均衡
					curr = DataSourceTypeConstant.READ + count.getAndAdd(1) % readDataSourceNum;
				}
				break;
			}
		}
		LOG.info(DataSourceLogConstant.CURR, curr);
		return curr;
	}

}
